import cv2,random
import numpy as np

v1 = np.load('train_data_v1.npy')
in_all = len(v1)
input('ALL IMGs:'+str(in_all))

w = [1,0,0,0,0,0,0]
s = [0,1,0,0,0,0,0]
d = [0,0,1,0,0,0,0]
a = [0,0,0,1,0,0,0]
skill = [0,0,0,0,1,0,0]
ak = [0,0,0,0,0,1,0]
big_skill = [0,0,0,0,0,0,1]
#w = np.array(w)
#s = np.array(s)
#d = np.array(d)
#a = np.array(a)
#skill = np.array(skill)
#ak = np.array(ak)
#big_skill = np.array(big_skill)
#stand = np.array(stand)

action_w = []
action_s = []
action_a = []
action_d = []
action_skill = []
action_ak = []
action_big = []


for data in v1:
    if w == data[1]:
        action_w.append(data)
    elif s == data[1]:
        action_s.append(data)
    elif d == data[1]:
        action_d.append(data)
    elif a == data[1]:
        action_a.append(data)
    elif skill == data[1]:
        action_skill.append(data)
    elif ak == data[1]:
        action_ak.append(data) 
    elif big_skill == data[1]:
        action_big.append(data)

        
#cut
a = int(len(action_a)*0.5)
action_a = action_a[a:]
a = int(len(action_w)*0.4)
action_w = action_w[a:]
a = int(len(action_d)*0.2)
action_d = action_d[a:]
a = int(len(action_s)*0.3)
action_s = action_s[a:]
a = int(len(action_skill)*0.7)
action_skill = action_skill[a:]
a = int(len(action_big)*0.8)
action_big = action_big[a:]
a = int(len(action_ak)*0.6)
action_ak = action_ak[a:]

v2 = action_a + action_ak + action_big + action_w + action_d + action_s + action_skill

random.shuffle(v2)

print(len(action_a))
print(len(action_w))
print(len(action_d))
print(len(action_s))
print(len(action_skill))
print(len(action_big))
print(len(action_ak))
print(len(v2))
input('ready!')

np.save('train_data_v2.npy',v2)
input('save!')
